package com.tianji.aigc.memory;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.json.JSONUtil;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.Message;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.stereotype.Component;

import java.util.List;

/**
 * Redis会话记忆处理类
 *
 * @Author mr.wu
 * @Date 2025-8-17 16:24
 */
@Component
public class RedisChatMemory implements ChatMemory {

    @Autowired
    private StringRedisTemplate redisTemplate;

    /**
     * 保存会话对应的消息到Redis中
     * @param conversationId
     * @param messages
     */
    @Override
    public void add(String conversationId, List<Message> messages) {
        if(CollUtil.isEmpty(messages)) {
            return;
        }
        //获取redisKey，组成结构：CHAT:用户id:sessionId
        String key = this.getKey(conversationId);
        //将Message消息转为JSON字符串
        List<String> messageList = messages.stream().map(message -> {
            return JSONUtil.toJsonStr(message);
        }).toList();
        //将消息JSON字符串列表添加到Redis的List中
        redisTemplate.opsForList().leftPushAll(key, messageList);
    }

    /**
     * 从Redis中查询某会话下的指定条数的消息列表
     * @param conversationId
     * @param lastN
     * @return
     */
    @Override
    public List<Message> get(String conversationId, int lastN) {
        if(lastN<=0){
            return List.of();
        }
        //获取redisKey，组成结构：CHAT:用户id:sessionId
        String key = this.getKey(conversationId);
        //从Redis的List中获取消息字符串列表
        List<String> messageStrList = redisTemplate.opsForList().range(key, 0, lastN);
        if(CollUtil.isEmpty(messageStrList)) {
            return List.of();
        }
        //将消息JSON字符串转为Message
        List<Message> messageList = messageStrList.stream().map(m -> {
            return JSONUtil.toBean(m, Message.class);
        }).toList();
        return messageList;
    }

    /**
     * 从Redis中删除某会话下的全部消息
     * @param conversationId
     */
    @Override
    public void clear(String conversationId) {
        String key = this.getKey(conversationId);
        redisTemplate.delete(key);
    }

    private String getKey(String conversationId) {
        return "CHAT:" + conversationId;
    }
}